from typing import AnyStr

import pandas as pd
import pymysql
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from config import *

engine = create_engine(f'mysql+pymysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_IP}:{MYSQL_PORT}/{MYSQL_QA_DB}')
session = sessionmaker(bind=engine)


def mysql_upload_data(data: pd.DataFrame):
    """
    一次性批量上传数据, 如果存在表格会覆盖
    :param data: 数据
    id: long
    question: str
    answer: str
    :return:
    """
    data.to_sql(name=MYSQL_QA_TABLE, con=engine, if_exists='replace', index=False)


def mysql_search_by_id(_id: str) -> AnyStr:
    with pymysql.connect(
            host=MYSQL_IP,
            port=MYSQL_PORT,
            user=MYSQL_USER,
            password=MYSQL_PASSWORD,
            database=MYSQL_QA_DB,
            charset='utf8'
    ) as conn:
        cursor = conn.cursor()
        sql = f"select answer from {MYSQL_QA_TABLE} where id='{_id}'"
        print(sql)
        cursor.execute(sql)
        mm = cursor.fetchone()
        res = mm[0] if mm else '数据库存储问题: 未找到答案'
        cursor.close()
        return res
